#include "crypto_helper.hpp"
#include "hypercube/hypercube_cryto_agent.hpp"
#include "hypercube/hypercube_hmm_engine.hpp"

#include "mem_usage.hpp"

//* true or false
#define FIX_NTL_RAND_SEED(enable)                                              \
  do {                                                                         \
    NTL::SetSeed(NTL::conv<NTL::ZZ>(enable ? 0l : std::time(nullptr)));        \
  } while (0)

#define RAND_NTL_MATRIX(mat, row, col, bl)                                     \
  do {                                                                         \
    mat.SetDims(row, col);                                                     \
    for (long i = 0; i < row; i++) {                                           \
      for (long j = 0; j < col; j++) {                                         \
        mat[i][j] = NTL::RandomBits_ZZ(bl);                                    \
      }                                                                        \
    }                                                                          \
  } while (0)

int main(int argc, char **argv) {
  bool sk;
  bool verbose;
  bool bsgs;
  bool single_thread;
  bool cache_free;
  bool minimal_m;

  std::size_t mem_begin_ = 0.0;
  std::size_t mem_end_ = 0.0;

  const double GB = static_cast<double>(1l << 30l);
  const double MB = static_cast<double>(1l << 20l);
  double mem_pre = 0.0;    // GB
  double mem_pk = 0.0;     // MB
  double mem_algor = 0.0;  // MB
  double mem_ctxt_A = 0.0; // MB
  double mem_ctxt_B = 0.0; // MB
  double mem_ctxt_C = 0.0; // MB

  unsigned long m, n, k;

  // Defualt BGV Params Set
  params<helib::BGV> params;

  // Create a `command_line_parser` using `params<helib::BGV>`
  command_line_parser parser(params);

  // Add Custom Command Line Options
  parser.add_option("-m", "4", "dim-m", &m);
  parser.add_option("-n", "4", "dim-n", &n);
  parser.add_option("-k", "4", "dim-k", &k);
  parser.add_option("-v", "verbos", &verbose);
  parser.add_option("-sk", "using_secret_key", &sk);
  parser.add_option("-bsgs", "bsgs", &bsgs);
  parser.add_option("-st", "single_thread", &single_thread);
  parser.add_option("-nocache", "always_build_context", &cache_free);
  parser.add_option("-minimal", "find_minimal_m", &minimal_m);

  //* Parse command line
  int ret = parser.parse(argc, argv);
  if (ret) {
    return ret;
  }

  // Select encryption method according to the command line option `-sk`
  auto key = sk ? ENC::SECKEY : ENC::PUBKEY;

  // Create a HMM status instance
  hmm_status<shmm_engine> hmm_status;
  // Select a HMM method according to the command line option `-bsgs`
  hmm_status.method(bsgs ? METHOD::BSGS : METHOD::NAIVE);
  // Set # of threads according to the command line option `-st`
  if (!single_thread) {
    hmm_status.threads(omp_get_max_threads());
  }

  // Get the excepted hypercube size
  // hmme.status()->eh_cube();
  const auto eh_cube = hypercube_hmme::get_hypercube(m, n, k);

  mem_begin_ = getCurrentRSS();
  // Initialize HE cryptograph system
  hypercube_hmmcc hmmcc(params, eh_cube, cache_free, minimal_m);
  mem_end_ = getCurrentRSS();
  mem_pre = static_cast<double>(mem_end_ - mem_begin_) / GB;

  std::stringstream ss;
  hmmcc.public_key().writeTo(ss);
  mem_pk = static_cast<double>(ss.str().size()) / MB;

  // Print out the HElib context
  hmmcc.report_context();

  // Create an inactive HMM engine
  hypercube_hmme hmme(hmm_status);
  //* Register the HMM engine using EA
  const auto eh_mnk = hypercube_hmme::get_expected_mnk(m, n, k);
  hmme.register_engine(hmmcc.context()->getEA(), eh_mnk);

  // Randomly generate input data
  NTL::mat_ZZ ptxt_A;
  NTL::mat_ZZ ptxt_B;
  FIX_NTL_RAND_SEED(false);
  RAND_NTL_MATRIX(ptxt_A, m, k, 8);
  RAND_NTL_MATRIX(ptxt_B, k, n, 8);

  // Encryption
  std::stringstream ss_A, ss_B;
  HELIB_NTIMER_START(_Encryption);
  auto ctxt_A = hmmcc.encrypt(ptxt_A, key);
  auto ctxt_B = hmmcc.encrypt(ptxt_B, key);
  HELIB_NTIMER_STOP(_Encryption);
  ctxt_A.writeTo(ss_A);
  ctxt_B.writeTo(ss_B);
  mem_ctxt_A = static_cast<double>(ss_A.str().size()) / MB;
  mem_ctxt_B = static_cast<double>(ss_B.str().size()) / MB;
  auto enc_time = _named_local_timer_Encryption.getTime() * 1000.0;
  auto init_cap = std::min(ctxt_A.capacity(), ctxt_B.capacity());

  // Execute HMM
  std::stringstream ss_C;
  auto ctxt_C = hmmcc.get_empty_ctxt();
  auto type = shmm_engine::parse_matrix_type(eh_mnk);
  mem_begin_ = getCurrentRSS();
  HELIB_NTIMER_START(_HMM);
  hmme.cgemm(type, ctxt_C, ctxt_A, ctxt_B);
  HELIB_NTIMER_STOP(_HMM);
  mem_end_ = getPeakRSS();
  mem_algor = static_cast<double>(mem_end_ - mem_begin_) / MB;

  ctxt_C.writeTo(ss_C);
  mem_ctxt_C = static_cast<double>(ss_C.str().size()) / MB;

  // Decryption
  HELIB_NTIMER_START(_Decryption);
  NTL::mat_ZZ C;
  hmmcc.decrypt(C, ctxt_C);
  HELIB_NTIMER_STOP(_Decryption);
  auto dec_time = _named_local_timer_Decryption.getTime() * 1000.0;
  if (verbose) {
    std::cout << "Dec(Enc(A) * Enc(B)):\n" << C << std::endl;
  }

  // Plaintext Verison
  if (verbose) {
    auto ptxt_AB = ptxt_A * ptxt_B;
    std::cout << "A * B:\n" << ptxt_AB << std::endl;
  }

  // Report HMM
  std::cout << "===================== HMM Report ======================\n";
  hmmcc.report_time();
  std::cout << "* Encryption Time            : " << enc_time << " ms\n";
  std::cout << "* Decryption Time            : " << dec_time << " ms\n";
  std::cout << "* Enc(A) Size                : " << mem_ctxt_A << " MB\n";
  std::cout << "* Enc(B) Size                : " << mem_ctxt_B << " MB\n";
  std::cout << "* Enc(C) Size                : " << mem_ctxt_C << " MB\n";
  std::cout << "* Context Mem Usage          : " << mem_pre << " GB\n";
  std::cout << "  * Public Key Size          : " << mem_pk << " MB\n";
  std::cout << "  * Other(pre-computation)   : " << mem_pre - (mem_pk / 1024.0)
            << " GB\n";
  std::cout << "* HMM Mem Usage              : " << mem_algor << " MB\n";
  std::cout << "* HMM Time (Outside)         : "
            << _named_local_timer_HMM.getTime() << " s\n";
  std::cout << "=======================================================\n";
  hmme.report();
  std::cout << "=======================================================\n";

  return 0;
}